package cn.xiaopengstack.config;

import org.springframework.ai.embedding.AbstractEmbeddingModel;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * @author jiangyangang
 */
@Configuration
@EnableConfigurationProperties(SpringAiProperties.class)
public class RAGEmbeddingConfig {

    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }


    @Bean
    public SimpleVectorStore vectorStore(OpenAiApi openAiApi) {
        AbstractEmbeddingModel embeddingModel = new OpenAiEmbeddingModel(openAiApi);
        return SimpleVectorStore.builder(embeddingModel).build();
    }

    @Bean
    public PgVectorStore pgVectorStore(SpringAiProperties properties, OpenAiApi openAiApi, JdbcTemplate jdbcTemplate) {
        EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(openAiApi);
        Integer dimensions = properties.getOpenAi().getEmbedding().getDimensions();

        return PgVectorStore.builder(jdbcTemplate, embeddingModel)
                .vectorTableName(properties.getPostgres().getTableName())
                .dimensions(dimensions)
                .build();
    }
}
